import xml.etree.ElementTree as ET


import numpy as np
import pycolmap

from pathlib import Path
import numpy as np
from scipy.spatial.transform import Rotation
import os

from pathlib import Path
from pprint import pformat

from hloc import (
    extract_features,
    match_features,
    pairs_from_covisibility,
    pairs_from_retrieval,
    pairs_from_poses,
)
from hloc import colmap_from_mvs, triangulation, localize_sfm, visualization

dataset = Path("/media/sdisk/vis_loc/220kv")  # change this if your dataset is somewhere else
images_path = dataset / "images-selected/"

outputs = Path("/media/sdisk/vis_loc/220kv/outputs/")  # where everything will be saved
sfm_pairs = outputs / "pairs-db-covis20.txt"  # top 20 most covisible in SIFT model
loc_pairs = outputs / "pairs-query-netvlad20.txt"  # top 20 retrieved by NetVLAD
reference_sfm = outputs / "sfm_superpoint+superglue"  # the SfM model we will build
results = outputs / "Aachen_hloc_superpoint+superglue_netvlad20.txt"  # the result file

# list the standard configurations available
print(f"Configs for feature extractors:\n{pformat(extract_features.confs)}")
print(f"Configs for feature matchers:\n{pformat(match_features.confs)}")


retrieval_conf = extract_features.confs["netvlad"]
feature_conf = extract_features.confs["superpoint_aachen"]
matcher_conf = match_features.confs["superglue"]

features = extract_features.main(feature_conf, images_path, outputs)

     
    
colmap_from_mvs.main(
    dataset / "mvs.xml",
    outputs / "sfm_sift",
)

#pairs_from_covisibility.main(outputs / "sfm_sift", sfm_pairs, num_matched=20)
pairs_from_poses.main(outputs / "sfm_sift", sfm_pairs, num_matched=20)

sfm_matches = match_features.main(
    matcher_conf, sfm_pairs, feature_conf["output"], outputs
)

reconstruction = triangulation.main(
    reference_sfm, outputs / "sfm_sift", images_path, sfm_pairs, features, sfm_matches
)

global_descriptors = extract_features.main(retrieval_conf, images_path, outputs)
# pairs_from_retrieval.main(
#      global_descriptors, loc_pairs, num_matched=20, db_prefix="db", query_prefix="query"
# )
